{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn \n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvLSTMCell(nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, hidden_dim, kernel_size, bias):\n",
    "        \"\"\"\n",
    "        Initialize ConvLSTM cell.\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_dim: int\n",
    "            Number of channels of input tensor.\n",
    "        hidden_dim: int\n",
    "            Number of channels of hidden state.\n",
    "        kernel_size: (int, int)\n",
    "            Size of the convolutional kernel.\n",
    "        bias: bool\n",
    "            Whether or not to add the bias.\n",
    "        \"\"\"\n",
    "\n",
    "        super(ConvLSTMCell, self).__init__()\n",
    "\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.kernel_size = kernel_size\n",
    "        self.padding = kernel_size[0] // 2, kernel_size[1] // 2\n",
    "        self.bias = bias\n",
    "\n",
    "        self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,\n",
    "                              out_channels=4 * self.hidden_dim,\n",
    "                              kernel_size=self.kernel_size,\n",
    "                              padding=self.padding,\n",
    "                              bias=self.bias).to(device)\n",
    "\n",
    "    def forward(self, input_tensor, cur_state):\n",
    "        h_cur, c_cur = cur_state\n",
    "\n",
    "        combined = torch.cat([input_tensor, h_cur], dim=1)  # concatenate along channel axis\n",
    "\n",
    "        combined_conv = self.conv(combined)\n",
    "        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)\n",
    "        i = torch.sigmoid(cc_i)\n",
    "        f = torch.sigmoid(cc_f)\n",
    "        o = torch.sigmoid(cc_o)\n",
    "        g = torch.tanh(cc_g)\n",
    "\n",
    "        c_next = f * c_cur + i * g\n",
    "        h_next = o * torch.tanh(c_next)\n",
    "\n",
    "        return h_next.to(device), c_next.to(device)\n",
    "\n",
    "    def init_hidden(self, batch_size, image_size):\n",
    "        height, width = image_size\n",
    "        return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),\n",
    "                torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))\n",
    "\n",
    "class ConvLSTM(nn.Module):\n",
    "\n",
    "    \"\"\"\n",
    "    Parameters:\n",
    "        input_dim: Number of channels in input\n",
    "        hidden_dim: Number of hidden channels\n",
    "        kernel_size: Size of kernel in convolutions\n",
    "        num_layers: Number of LSTM layers stacked on each other\n",
    "        batch_first: Whether or not dimension 0 is the batch or not\n",
    "        bias: Bias or no bias in Convolution\n",
    "        return_all_layers: Return the list of computations for all layers\n",
    "        Note: Will do same padding.\n",
    "    Input:\n",
    "        A tensor of size B, T, C, H, W or T, B, C, H, W\n",
    "    Output:\n",
    "        A tuple of two lists of length num_layers (or length 1 if return_all_layers is False).\n",
    "            0 - layer_output_list is the list of lists of length T of each output\n",
    "            1 - last_state_list is the list of last states\n",
    "                    each element of the list is a tuple (h, c) for hidden state and memory\n",
    "    Example:\n",
    "        >> x = torch.rand((32, 10, 64, 128, 128))\n",
    "        >> convlstm = ConvLSTM(64, 16, 3, 1, True, True, False)\n",
    "        >> _, last_states = convlstm(x)\n",
    "        >> h = last_states[0][0]  # 0 for layer index, 0 for h index\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_dim, hidden_dim, kernel_size, num_layers,\n",
    "                 batch_first=False, bias=True, return_all_layers=False):\n",
    "        super(ConvLSTM, self).__init__()\n",
    "\n",
    "        #self._check_kernel_size_consistency(kernel_size)\n",
    "\n",
    "        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers\n",
    "        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)\n",
    "        hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)\n",
    "        if not len(kernel_size) == len(hidden_dim) == num_layers:\n",
    "            raise ValueError('Inconsistent list length.')\n",
    "\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.kernel_size = kernel_size\n",
    "        self.num_layers = num_layers\n",
    "        self.batch_first = batch_first\n",
    "        self.bias = bias\n",
    "        self.return_all_layers = return_all_layers\n",
    "\n",
    "        cell_list = []\n",
    "        for i in range(0, self.num_layers):\n",
    "            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]\n",
    "            cell_list.append(ConvLSTMCell(input_dim=cur_input_dim,\n",
    "                                          hidden_dim=self.hidden_dim[i],\n",
    "                                          kernel_size=self.kernel_size[i],\n",
    "                                          bias=self.bias))\n",
    "\n",
    "        self.cell_list = nn.ModuleList(cell_list)\n",
    "\n",
    "        self.liner = nn.Linear(16*24*72,24)\n",
    "    def forward(self, input_tensor, hidden_state=None):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_tensor: todo\n",
    "            5-D Tensor either of shape (t, b, c, h, w) or (b, t, c, h, w)\n",
    "        hidden_state: todo\n",
    "            None. todo implement stateful\n",
    "        Returns\n",
    "        -------\n",
    "        last_state_list, layer_output\n",
    "        \"\"\"\n",
    "        if not self.batch_first:\n",
    "            # (t, b, c, h, w) -> (b, t, c, h, w)\n",
    "            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)\n",
    "\n",
    "        b, _, _, h, w = input_tensor.size()\n",
    "        # Implement stateful ConvLSTM\n",
    "        if hidden_state is not None:\n",
    "            raise NotImplementedError()\n",
    "        else:\n",
    "            # Since the init is done in forward. Can send image size here\n",
    "            hidden_state = self._init_hidden(batch_size=b,\n",
    "                                             image_size=(h, w))\n",
    "\n",
    "        layer_output_list = []\n",
    "        last_state_list = []\n",
    "\n",
    "        seq_len = input_tensor.size(1)\n",
    "        cur_layer_input = input_tensor\n",
    "\n",
    "        for layer_idx in range(self.num_layers):\n",
    "            h, c = hidden_state[layer_idx]\n",
    "            output_inner = []\n",
    "            for t in range(seq_len):\n",
    "                h, c = self.cell_list[layer_idx](input_tensor=cur_layer_input[:, t, :, :, :],\n",
    "                                                 cur_state=[h, c])\n",
    "                output_inner.append(h)\n",
    "\n",
    "            layer_output = torch.stack(output_inner, dim=1)\n",
    "            cur_layer_input = layer_output\n",
    "\n",
    "            layer_output_list.append(layer_output)\n",
    "            last_state_list.append([h, c])\n",
    "        if not self.return_all_layers:\n",
    "            layer_output_list = layer_output_list[-1:]\n",
    "            last_state_list = last_state_list[-1:]\n",
    "\n",
    "        return layer_output_list, last_state_list\n",
    "\n",
    "    def _init_hidden(self, batch_size, image_size):\n",
    "        init_states = []\n",
    "        for i in range(self.num_layers):\n",
    "            init_states.append(self.cell_list[i].init_hidden(batch_size, image_size))\n",
    "        return init_states\n",
    "\n",
    "    @staticmethod\n",
    "    def _check_kernel_size_consistency(kernel_size):\n",
    "        if not (isinstance(kernel_size, tuple) or\n",
    "                (isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):\n",
    "            raise ValueError('`kernel_size` must be tuple or list of tuples')\n",
    "\n",
    "    @staticmethod\n",
    "    def _extend_for_multilayer(param, num_layers):\n",
    "        if not isinstance(param, list):\n",
    "            param = [param] * num_layers\n",
    "        return param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "h_state: torch.Size([32, 16, 128, 128])\n",
      "input: torch.Size([32, 10, 64, 128, 128])\n",
      "out: torch.Size([32, 10, 16, 128, 128])\n"
     ]
    }
   ],
   "source": [
    "device = 'cpu'\n",
    "#(batch  seq_len  channel h,w)\n",
    "x = torch.rand((32, 10, 64, 128, 128))\n",
    "y = torch.rand((32, 10, 16, 128, 128))\n",
    "\n",
    "\n",
    "# input_channel=64  hidden_dim=16  ker=3 ,numlayers = 1   batchfirst =True  bias=True  return_all_layers=False\n",
    "convlstm = ConvLSTM(64, 16, (3,3), 1, True, True, False)\n",
    "out, last_states = convlstm(x)\n",
    "\n",
    "# out (32, 10, 16, 128, 128)  (batch  seq_len  channel h,w)   \n",
    "# last_states(32,16,128,128)\n",
    "h = last_states[0][0]  # 0 for layer index, 0 for h index   h0 c0\n",
    "print('h_state:',h.shape)\n",
    "print('input:',x.shape)\n",
    "print('out:',out[0].shape)  #0 for layer index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Instantiate model\n",
      "ConvLSTM(\n",
      "  (cell_list): ModuleList(\n",
      "    (0): ConvLSTMCell(\n",
      "      (conv): Conv2d(28, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    )\n",
      "  )\n",
      "  (liner): Linear(in_features=27648, out_features=24, bias=True)\n",
      ")\n",
      "Create input and target Variables\n",
      "Create a MSE criterion\n",
      "Run for 2 iterations\n",
      "torch.Size([32, 10, 16, 24, 72])\n",
      " > Epoch  1 loss: 0.325\n",
      "torch.Size([32, 10, 16, 24, 72])\n",
      " > Epoch  2 loss: 0.324\n",
      "Input size: [32, 10, 12, 24, 72]\n",
      "Target size: [32, 10, 16, 24, 72]\n",
      "Last hidden state size: [32, 16, 24, 72]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "lr = 1e-1       # learning rate\n",
    "max_epoch = 2  # number of epochs\n",
    "\n",
    "# set manual seed\n",
    "torch.manual_seed(0)\n",
    "\n",
    "print('Instantiate model')\n",
    "device = 'cpu'\n",
    "# input_channel=64  hidden_dim=16  ker=3 ,numlayers = 1   batchfirst =True  bias=True  return_all_layers=False\n",
    "model = ConvLSTM(12, 16, (3,3), 1, True, True, False)\n",
    "print(repr(model))\n",
    "\n",
    "print('Create input and target Variables')\n",
    "# x = batch  seq_len  channel h,w\n",
    "x = torch.rand((32, 10, 12, 24, 72))\n",
    "y = torch.rand((32,10,16,24,72))\n",
    "\n",
    "\n",
    "print('Create a MSE criterion')\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.00001)\n",
    "\n",
    "print('Run for', max_epoch, 'iterations')\n",
    "for epoch in range(0, max_epoch):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    out,h = model(x)\n",
    "    print(out[0].shape)\n",
    "    loss = loss_fn(out[0],y)\n",
    "    print(' > Epoch {:2d} loss: {:.3f}'.format((epoch+1), loss.item()))\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "print('Input size:', list(x.size()))\n",
    "print('Target size:', list(y.size()))\n",
    "print('Last hidden state size:', list(h[0][0].size()))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
